#include <windows.h>
#include <iostream>
#include <functional>
#include <thread>
#include <atomic>
#include <string>
#include <vector>
#include <map>
#include <TlHelp32.h>
#include <psapi.h>
#include <mutex>
#include <memory>
#include <unordered_set>
#include <fcntl.h>
#include <io.h>
#include <commctrl.h>

using namespace std;

// 全局钩子DLL管理类
class GlobalHookManager {
private:
    HMODULE hHookDll = nullptr;
    function<BOOL()> InstallGlobalHooks = nullptr;
    function<BOOL()> UninstallGlobalHooks = nullptr;
    function<void(DWORD)> SetMainProcessId = nullptr;

public:
    bool Load() {
        wcout << L"Loading GlobalHooks.dll..." << endl;
        hHookDll = LoadLibraryW(L"GlobalHooks.dll");
        if (!hHookDll) {
            DWORD error = GetLastError();
            wcerr << L"Failed to load GlobalHooks.dll: " << error << endl;
            return false;
        }
        wcout << L"GlobalHooks.dll loaded successfully" << endl;

        InstallGlobalHooks = reinterpret_cast<BOOL(*)()>(GetProcAddress(hHookDll, "InstallGlobalHooks"));
        UninstallGlobalHooks = reinterpret_cast<BOOL(*)()>(GetProcAddress(hHookDll, "UninstallGlobalHooks"));
        SetMainProcessId = reinterpret_cast<void(*)(DWORD)>(GetProcAddress(hHookDll, "SetMainProcessId"));

        if (!InstallGlobalHooks) {
            wcerr << L"Failed to get InstallGlobalHooks function" << endl;
        }
        if (!UninstallGlobalHooks) {
            wcerr << L"Failed to get UninstallGlobalHooks function" << endl;
        }
        if (!SetMainProcessId) {
            wcerr << L"Failed to get SetMainProcessId function" << endl;
        }

        if (!InstallGlobalHooks || !UninstallGlobalHooks || !SetMainProcessId) {
            wcerr << L"Failed to get function pointers from DLL" << endl;
            FreeLibrary(hHookDll);
            hHookDll = nullptr;
            return false;
        }

        wcout << L"Setting main process ID: " << GetCurrentProcessId() << endl;
        SetMainProcessId(GetCurrentProcessId());
        return true;
    }

    bool Install() {
        wcout << L"Installing global hooks..." << endl;
        if (InstallGlobalHooks && InstallGlobalHooks()) {
            wcout << L"Global hooks installed successfully" << endl;
            return true;
        }
        wcerr << L"Failed to install global hooks" << endl;
        return false;
    }

    void Uninstall() {
        wcout << L"Uninstalling global hooks..." << endl;
        if (UninstallGlobalHooks) {
            UninstallGlobalHooks();
        }
        if (hHookDll) {
            wcout << L"Freeing GlobalHooks.dll" << endl;
            FreeLibrary(hHookDll);
            hHookDll = nullptr;
        }
    }

    ~GlobalHookManager() {
        Uninstall();
    }
};

class WindowManager {
public:
    struct TargetInfo {
        HWND hHost;
        LONG_PTR originalStyle;
        LONG_PTR originalExStyle;
        RECT originalRect;
    };

    static atomic<bool> classRegistered;
    static mutex classMutex;

    bool IsRunning() const { return running.load(); }

    static vector<HWND> GetProcessWindows(DWORD pid) {
        vector<HWND> windows;

        struct EnumData {
            DWORD pid;
            vector<HWND>* windows;
        } data = { pid, &windows };

        auto enumProc = [](HWND hwnd, LPARAM lParam) -> BOOL {
            EnumData* data = reinterpret_cast<EnumData*>(lParam);
            DWORD windowPid;
            GetWindowThreadProcessId(hwnd, &windowPid);

            if (windowPid == data->pid &&
                IsWindowVisible(hwnd) &&
                GetParent(hwnd) == NULL) {
                wchar_t title[256];
                GetWindowTextW(hwnd, title, 256);
                wcout << L"Found target window: " << hwnd << L" Title: " << title << endl;
                data->windows->push_back(hwnd);
            }
            return TRUE;
            };

        EnumWindows(enumProc, reinterpret_cast<LPARAM>(&data));
        return windows;
    }

    void AddTargetWindow(HWND hTarget) {
        {
            lock_guard<mutex> lock(windowsMutex);
            if (hostWindows.find(hTarget) != hostWindows.end()) {
                wcout << L"Window already managed: " << hTarget << endl;
                return;
            }
        }

        wchar_t title[256];
        if (GetWindowTextW(hTarget, title, 256)) {
            wcout << L"Adding target window: " << hTarget << L" Title: " << title << endl;
        }
        else {
            wcout << L"Adding target window: " << hTarget << L" (no title)" << endl;
        }

        thread t(&WindowManager::WindowManagementThread, this, hTarget);
        t.detach();
    }

    void Stop() {
        running = false;

        vector<HWND> hosts;
        {
            lock_guard<mutex> lock(windowsMutex);
            for (auto& pair : hostWindows) {
                hosts.push_back(pair.second.hHost);
            }
        }

        for (HWND hHost : hosts) {
            if (IsWindow(hHost)) {
                wcout << L"Closing host window: " << hHost << endl;

                // 设置标志表示正常关闭
                SetPropW(hHost, L"MANAGED_CLOSE", (HANDLE)TRUE);
                PostMessage(hHost, WM_CLOSE, 0, 0);
            }
        }

        Sleep(1000);

        lock_guard<mutex> lock(windowsMutex);
        hostWindows.clear();
    }

private:
    map<HWND, TargetInfo> hostWindows;
    atomic<bool> running{ true };
    mutex windowsMutex;

    static LRESULT CALLBACK HostWndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

    HWND CreateHostWindow(HWND hTarget) {
        HINSTANCE hInstance = GetModuleHandle(NULL);
        const wchar_t* className = L"HostWindowClass";

        // 确保窗口类只注册一次
        {
            lock_guard<mutex> lock(classMutex);
            if (!classRegistered) {
                WNDCLASSEXW wc = { sizeof(WNDCLASSEX) };
                wc.lpfnWndProc = HostWndProc;
                wc.hInstance = hInstance;
                wc.lpszClassName = className;
                wc.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);

                if (RegisterClassExW(&wc)) {
                    classRegistered = true;
                    wcout << L"Window class registered successfully." << endl;
                }
                else {
                    DWORD error = GetLastError();
                    // 1410 表示已经注册过，可以接受
                    if (error != ERROR_CLASS_ALREADY_EXISTS) {
                        wcerr << L"RegisterClassEx failed: " << error << endl;
                        return NULL;
                    }
                    classRegistered = true;
                    wcout << L"Window class already exists, continuing..." << endl;
                }
            }
        }

        wchar_t title[256];
        wstring windowTitle = L"Protected Host Window";
        if (GetWindowTextW(hTarget, title, 256)) {
            windowTitle += L" for: ";
            windowTitle += title;
        }

        HWND hHost = CreateWindowExW(
            0, className, windowTitle.c_str(),
            WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, CW_USEDEFAULT,
            800, 600, NULL, NULL, hInstance, NULL
        );
        if (hHost) {
            SetWindowLongPtr(hHost, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(this));
            wcout << L"Created host window: " << hHost << L" for target: " << hTarget << endl;

            // 禁用系统菜单中的关闭按钮
            HMENU hMenu = GetSystemMenu(hHost, FALSE);
            if (hMenu) {
                EnableMenuItem(hMenu, SC_CLOSE, MF_BYCOMMAND | MF_GRAYED);
                DrawMenuBar(hHost);
                wcout << L"Disabled close button for host window: " << hHost << endl;
            }
        }
        else {
            wcerr << L"CreateWindowEx failed: " << GetLastError() << endl;
        }

        return hHost;
    }

    void EmbedWindow(HWND hTarget, HWND hHost) {
        if (!IsWindow(hTarget)) {
            wcerr << L"EmbedWindow: Target window invalid: " << hTarget << endl;
            return;
        }
        if (!IsWindow(hHost)) {
            wcerr << L"EmbedWindow: Host window invalid: " << hHost << endl;
            return;
        }

        // 保存原始状态
        TargetInfo info;
        info.hHost = hHost;
        info.originalStyle = GetWindowLongPtr(hTarget, GWL_STYLE);
        info.originalExStyle = GetWindowLongPtr(hTarget, GWL_EXSTYLE);
        GetWindowRect(hTarget, &info.originalRect);

        // 设置窗口属性以存储宿主窗口句柄
        if (!SetPropW(hTarget, L"HOST_WINDOW", hHost)) {
            DWORD error = GetLastError();
            wcerr << L"SetPropW failed: " << error << endl;
        }
        else {
            wcout << L"Set HOST_WINDOW property on target window: " << hTarget
                << L" to host: " << hHost << endl;
        }

        {
            lock_guard<mutex> lock(windowsMutex);
            hostWindows[hTarget] = info;
        }

        // 修改窗口样式
        LONG_PTR newStyle = info.originalStyle;
        newStyle &= ~WS_POPUP;
        newStyle |= WS_CHILD;

        if (SetWindowLongPtr(hTarget, GWL_STYLE, newStyle) == 0) {
            wcerr << L"SetWindowLongPtr(GWL_STYLE) failed: " << GetLastError() << endl;
        }

        // 设置父子关系
        if (!SetParent(hTarget, hHost)) {
            wcerr << L"SetParent failed: " << GetLastError() << endl;
        }

        // 调整尺寸
        RECT rect;
        if (!GetClientRect(hHost, &rect)) {
            wcerr << L"GetClientRect failed: " << GetLastError() << endl;
            return;
        }

        if (!SetWindowPos(hTarget, NULL, 0, 0, rect.right, rect.bottom,
            SWP_NOZORDER | SWP_FRAMECHANGED)) {
            wcerr << L"SetWindowPos failed: " << GetLastError() << endl;
        }

        // 显示窗口
        ShowWindow(hHost, SW_SHOW);
        UpdateWindow(hHost);
        wcout << L"Successfully embedded window: " << hTarget << L" into host: " << hHost << endl;
    }

    void RestoreTargetWindow(HWND hTarget) {
        if (!IsWindow(hTarget)) return;

        lock_guard<mutex> lock(windowsMutex);
        auto it = hostWindows.find(hTarget);
        if (it != hostWindows.end()) {
            auto& info = it->second;

            // 恢复原始样式
            SetWindowLongPtr(hTarget, GWL_STYLE, info.originalStyle);
            SetWindowLongPtr(hTarget, GWL_EXSTYLE, info.originalExStyle);
            SetParent(hTarget, NULL);

            // 恢复原始位置和大小
            SetWindowPos(hTarget, NULL,
                info.originalRect.left, info.originalRect.top,
                info.originalRect.right - info.originalRect.left,
                info.originalRect.bottom - info.originalRect.top,
                SWP_NOZORDER | SWP_FRAMECHANGED);

            // 移除属性
            RemovePropW(hTarget, L"HOST_WINDOW");

            hostWindows.erase(it);
            wcout << L"Restored target window: " << hTarget << endl;
        }
    }

    void WindowManagementThread(HWND hTarget) {
        wcout << L"Starting management thread for window: " << hTarget << endl;

        HWND hHost = CreateHostWindow(hTarget);
        if (!hHost) {
            wcerr << L"CreateHostWindow failed for target: " << hTarget << endl;
            return;
        }

        EmbedWindow(hTarget, hHost);

        MSG msg;
        while (GetMessage(&msg, NULL, 0, 0)) {
            TranslateMessage(&msg);
            DispatchMessage(&msg);

            if (!IsWindow(hTarget)) {
                wcout << L"Target window destroyed: " << hTarget << endl;
                if (IsWindow(hHost)) {
                    // 设置标志表示正常关闭
                    SetPropW(hHost, L"MANAGED_CLOSE", (HANDLE)TRUE);
                    DestroyWindow(hHost);
                }
                break;
            }

            // 定期检查父子关系
            if (GetParent(hTarget) != hHost) {
                wcout << L"Re-parenting target window: " << hTarget << endl;
                SetParent(hTarget, hHost);

                RECT rect;
                if (GetClientRect(hHost, &rect)) {
                    SetWindowPos(hTarget, NULL, 0, 0, rect.right, rect.bottom,
                        SWP_NOZORDER | SWP_FRAMECHANGED);
                }
            }
        }

        // 恢复目标窗口
        RestoreTargetWindow(hTarget);

        wcout << L"Exiting management thread for window: " << hTarget << endl;
    }
};

// 静态成员初始化
atomic<bool> WindowManager::classRegistered(false);
mutex WindowManager::classMutex;

// 宿主窗口过程
LRESULT CALLBACK WindowManager::HostWndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    WindowManager* manager = reinterpret_cast<WindowManager*>(GetWindowLongPtrW(hWnd, GWLP_USERDATA));
    if (!manager) return DefWindowProcW(hWnd, msg, wParam, lParam);

    HWND hTarget = nullptr;
    {
        lock_guard<mutex> lock(manager->windowsMutex);
        for (const auto& pair : manager->hostWindows) {
            if (pair.second.hHost == hWnd) {
                hTarget = pair.first;
                break;
            }
        }
    }

    switch (msg) {
    case WM_SIZE: {
        if (IsWindow(hTarget)) {
            RECT rect;
            GetClientRect(hWnd, &rect);
            SetWindowPos(hTarget, NULL, 0, 0, rect.right, rect.bottom,
                SWP_NOZORDER | SWP_FRAMECHANGED);
        }
        break;
    }
    case WM_KEYDOWN: {
        if (wParam == VK_ESCAPE) { // 检测ESC键
            // 使用窗口属性存储全屏状态
            bool isFullscreen = GetPropW(hWnd, L"IS_FULLSCREEN") != NULL;

            if (!isFullscreen) {
                // 进入全屏模式
                SetPropW(hWnd, L"IS_FULLSCREEN", (HANDLE)TRUE);

                // 保存当前窗口位置和样式
                RECT originalRect;
                GetWindowRect(hWnd, &originalRect);
                SetPropW(hWnd, L"ORIGINAL_RECT", (HANDLE)new RECT(originalRect));

                LONG_PTR originalStyle = GetWindowLongPtrW(hWnd, GWL_STYLE);
                SetPropW(hWnd, L"ORIGINAL_STYLE", (HANDLE)originalStyle);

                LONG_PTR originalExStyle = GetWindowLongPtrW(hWnd, GWL_EXSTYLE);
                SetPropW(hWnd, L"ORIGINAL_EXSTYLE", (HANDLE)originalExStyle);

                // 获取屏幕尺寸
                int screenWidth = GetSystemMetrics(SM_CXSCREEN);
                int screenHeight = GetSystemMetrics(SM_CYSCREEN);

                // 设置全屏样式
                SetWindowLongPtrW(hWnd, GWL_STYLE,
                    originalStyle & ~(WS_CAPTION | WS_THICKFRAME));
                SetWindowLongPtrW(hWnd, GWL_EXSTYLE,
                    originalExStyle & ~(WS_EX_DLGMODALFRAME |
                        WS_EX_WINDOWEDGE | WS_EX_CLIENTEDGE | WS_EX_STATICEDGE));

                // 调整窗口到全屏
                SetWindowPos(hWnd, HWND_TOP, 0, 0, screenWidth, screenHeight,
                    SWP_FRAMECHANGED | SWP_SHOWWINDOW);

                wcout << L"Entered fullscreen mode for window: " << hWnd << endl;
            }
            else {
                // 退出全屏模式
                SetPropW(hWnd, L"IS_FULLSCREEN", (HANDLE)FALSE);

                // 恢复原始样式
                LONG_PTR originalStyle = (LONG_PTR)GetPropW(hWnd, L"ORIGINAL_STYLE");
                SetWindowLongPtrW(hWnd, GWL_STYLE, originalStyle);

                LONG_PTR originalExStyle = (LONG_PTR)GetPropW(hWnd, L"ORIGINAL_EXSTYLE");
                SetWindowLongPtrW(hWnd, GWL_EXSTYLE, originalExStyle);

                // 恢复原始位置和大小
                RECT* originalRect = (RECT*)GetPropW(hWnd, L"ORIGINAL_RECT");
                SetWindowPos(hWnd, NULL,
                    originalRect->left, originalRect->top,
                    originalRect->right - originalRect->left,
                    originalRect->bottom - originalRect->top,
                    SWP_FRAMECHANGED | SWP_SHOWWINDOW);

                // 清理资源
                delete originalRect;
                RemovePropW(hWnd, L"ORIGINAL_RECT");
                RemovePropW(hWnd, L"ORIGINAL_STYLE");
                RemovePropW(hWnd, L"ORIGINAL_EXSTYLE");

                wcout << L"Exited fullscreen mode for window: " << hWnd << endl;
            }
        }
        break;
    }
    case WM_CLOSE: {
        // 检查是否是我们自己发起的关闭
        if (GetPropW(hWnd, L"MANAGED_CLOSE")) {
            RemovePropW(hWnd, L"MANAGED_CLOSE");
            DestroyWindow(hWnd);
        }
        else {
            wcout << L"Preventing window close: " << hWnd << endl;
        }
        return 0;
    }
    case WM_DESTROY: {
        wcout << L"Destroying host window: " << hWnd << endl;

        // 恢复目标窗口
        if (hTarget) {
            manager->RestoreTargetWindow(hTarget);
        }

        // 清理全屏相关属性
        if (GetPropW(hWnd, L"ORIGINAL_RECT")) {
            delete (RECT*)GetPropW(hWnd, L"ORIGINAL_RECT");
            RemovePropW(hWnd, L"ORIGINAL_RECT");
        }
        RemovePropW(hWnd, L"ORIGINAL_STYLE");
        RemovePropW(hWnd, L"ORIGINAL_EXSTYLE");
        RemovePropW(hWnd, L"IS_FULLSCREEN");

        PostQuitMessage(0);
        break;
    }
    case WM_SYSCOMMAND:
        // 防止通过Alt+F4关闭
        if ((wParam & 0xFFF0) == SC_CLOSE) {
            wcout << L"Blocked Alt+F4 for window: " << hWnd << endl;
            return 0;
        }
        break;
    }
    return DefWindowProcW(hWnd, msg, wParam, lParam);
}

class ProcessMonitor {
private:
    DWORD processId;
    HANDLE processHandle = NULL;
    WindowManager windowManager;
    atomic<bool> isMonitoring{ false };
    thread monitorThread;

    void monitorLoop() {
        const DWORD checkInterval = 5000;
        DWORD lastCheck = GetTickCount();

        while (isMonitoring && windowManager.IsRunning()) {
            if (isRunning()) {
                DWORD currentTime = GetTickCount();
                if (currentTime - lastCheck >= checkInterval) {
                    lastCheck = currentTime;

                    wcout << L"Checking windows for PID: " << processId << endl;
                    vector<HWND> windows = WindowManager::GetProcessWindows(processId);
                    wcout << L"Found " << windows.size() << L" windows for PID: " << processId << endl;

                    for (HWND hWnd : windows) {
                        windowManager.AddTargetWindow(hWnd);
                    }
                }
            }
            Sleep(100);
        }
    }

public:
    using ExitCallback = std::function<void(DWORD)>;
    ExitCallback callbackFunc;

    ProcessMonitor(DWORD pid, ExitCallback callback = nullptr)
        : processId(pid), callbackFunc(callback) {
    }

    ~ProcessMonitor() {
        stop();
    }

    bool start() {
        if (isMonitoring) {
            return false;
        }

        processHandle = OpenProcess(PROCESS_QUERY_INFORMATION | SYNCHRONIZE, FALSE, processId);
        if (processHandle == NULL) {
            wcerr << L"OpenProcess failed: " << GetLastError() << endl;
            return false;
        }

        isMonitoring = true;
        monitorThread = thread(&ProcessMonitor::monitorLoop, this);
        wcout << L"Started monitoring for PID: " << processId << endl;
        return true;
    }

    void stop() {
        if (isMonitoring) {
            wcout << L"Stopping monitoring for PID: " << processId << endl;
            isMonitoring = false;
            if (monitorThread.joinable()) {
                monitorThread.join();
            }
        }

        windowManager.Stop();

        if (processHandle != NULL) {
            CloseHandle(processHandle);
            processHandle = NULL;
        }
    }

    bool isRunning() {
        if (processHandle == NULL) return false;

        DWORD exitCode;
        if (GetExitCodeProcess(processHandle, &exitCode)) {
            return exitCode == STILL_ACTIVE;
        }
        return false;
    }
};

// 获取指定进程名的所有PID（包括路径匹配）
vector<DWORD> GetPIDsByProcessNameWithPath(const wstring& processName) {
    vector<DWORD> pids;

    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnapshot == INVALID_HANDLE_VALUE) {
        wcerr << L"CreateToolhelp32Snapshot failed: " << GetLastError() << endl;
        return pids;
    }

    PROCESSENTRY32W pe;
    pe.dwSize = sizeof(PROCESSENTRY32W);

    if (!Process32FirstW(hSnapshot, &pe)) {
        wcerr << L"Process32First failed: " << GetLastError() << endl;
        CloseHandle(hSnapshot);
        return pids;
    }

    do {
        wstring exeFile = pe.szExeFile;
        // 检查文件名是否匹配（不区分大小写）
        if (_wcsicmp(exeFile.c_str(), processName.c_str()) == 0) {
            pids.push_back(pe.th32ProcessID);
            wcout << L"Found process: " << processName << L" PID: " << pe.th32ProcessID << endl;
        }
        // 检查是否包含完整路径
        else {
            wchar_t fullPath[MAX_PATH];
            HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pe.th32ProcessID);
            if (hProcess) {
                DWORD size = MAX_PATH;
                if (QueryFullProcessImageNameW(hProcess, 0, fullPath, &size)) {
                    wstring path(fullPath);
                    // 提取文件名部分
                    size_t pos = path.find_last_of(L"\\/");
                    wstring fileName = (pos != wstring::npos) ? path.substr(pos + 1) : path;

                    if (_wcsicmp(fileName.c_str(), processName.c_str()) == 0) {
                        pids.push_back(pe.th32ProcessID);
                        wcout << L"Found process by path: " << fileName << L" PID: " << pe.th32ProcessID
                            << L" Path: " << path << endl;
                    }
                }
                CloseHandle(hProcess);
            }
        }
    } while (Process32NextW(hSnapshot, &pe));

    CloseHandle(hSnapshot);
    return pids;
}

// DLL注入函数
void InjectDllIntoProcess(DWORD pid, const wchar_t* dllPath) {
    wcout << L"Attempting to inject DLL into PID: " << pid << endl;
    wcout << L"DLL Path: " << dllPath << endl;

    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
    if (!hProcess) {
        wcerr << L"OpenProcess failed: " << GetLastError() << endl;
        return;
    }

    // 在目标进程中分配内存
    size_t pathSize = (wcslen(dllPath) + 1) * sizeof(wchar_t);
    LPVOID pDllPath = VirtualAllocEx(hProcess, NULL, pathSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (!pDllPath) {
        wcerr << L"VirtualAllocEx failed: " << GetLastError() << endl;
        CloseHandle(hProcess);
        return;
    }

    // 写入DLL路径
    if (!WriteProcessMemory(hProcess, pDllPath, dllPath, pathSize, NULL)) {
        wcerr << L"WriteProcessMemory failed: " << GetLastError() << endl;
        VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return;
    }

    // 获取LoadLibraryW地址
    LPTHREAD_START_ROUTINE pLoadLibrary = (LPTHREAD_START_ROUTINE)
        GetProcAddress(GetModuleHandleW(L"kernel32.dll"), "LoadLibraryW");
    if (!pLoadLibrary) {
        wcerr << L"GetProcAddress failed: " << GetLastError() << endl;
        VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return;
    }

    // 创建远程线程加载DLL
    HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0, pLoadLibrary, pDllPath, 0, NULL);
    if (!hThread) {
        wcerr << L"CreateRemoteThread failed: " << GetLastError() << endl;
        VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return;
    }

    // 等待线程完成
    WaitForSingleObject(hThread, INFINITE);

    // 检查注入结果
    DWORD exitCode;
    GetExitCodeThread(hThread, &exitCode);
    if (exitCode == 0) {
        wcerr << L"DLL injection failed in target process" << endl;
    }
    else {
        wcout << L"DLL injected successfully into PID: " << pid << endl;
    }

    // 清理
    CloseHandle(hThread);
    VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
    CloseHandle(hProcess);
}

int main() {
    _setmode(_fileno(stdout), _O_U16TEXT); // 确保控制台输出Unicode

    wcout << L"Starting protected window manager with global hooks..." << endl;

    // 加载全局钩子DLL到主进程
    GlobalHookManager hookManager;
    if (!hookManager.Load() || !hookManager.Install()) {
        wcerr << L"Critical error: Failed to initialize global hooks" << endl;
        return 1;
    }

    vector<unique_ptr<ProcessMonitor>> monitors;
    unordered_set<DWORD> monitoredPIDs; // 用于跟踪已监控的PID，避免重复

    wcout << L"Enter target process names (one per line, enter empty line to finish):" << endl;
    wstring processName;

    while (true) {
        wcout << L"Process name: ";
        getline(wcin, processName);

        if (processName.empty()) {
            break;
        }

        // 获取匹配的所有PID
        vector<DWORD> pids = GetPIDsByProcessNameWithPath(processName);

        if (pids.empty()) {
            wcout << L"No running processes found for: " << processName << endl;
            continue;
        }

        // 获取当前DLL路径
        wchar_t dllPath[MAX_PATH];
        GetModuleFileNameW(NULL, dllPath, MAX_PATH);
        wchar_t* lastSlash = wcsrchr(dllPath, L'\\');
        if (lastSlash) {
            wcscpy_s(lastSlash + 1, MAX_PATH - (lastSlash - dllPath) - 1, L"GlobalHooks.dll");
        }

        for (DWORD pid : pids) {
            // 检查是否已经监控过这个PID
            if (monitoredPIDs.find(pid) != monitoredPIDs.end()) {
                wcout << L"PID " << pid << L" is already being monitored." << endl;
                continue;
            }

            // 注入DLL到目标进程
            InjectDllIntoProcess(pid, dllPath);

            auto monitor = make_unique<ProcessMonitor>(pid, [pid](DWORD exitCode) {
                wcout << L"Process " << pid << L" exited with code: " << exitCode << endl;
                });

            if (monitor->start()) {
                monitors.push_back(std::move(monitor));
                monitoredPIDs.insert(pid);
                wcout << L"Monitoring started for PID: " << pid << endl;
            }
            else {
                wcerr << L"Failed to start monitoring for PID: " << pid << endl;
            }
        }
    }

    if (monitors.empty()) {
        wcerr << L"No valid processes to monitor." << endl;
        return 1;
    }

    wcout << L"Host windows are protected from accidental closure." << endl;
    wcout << L"Press Enter to stop monitoring..." << endl;
    cin.get();

    // 停止所有监控
    for (auto& monitor : monitors) {
        monitor->stop();
    }

    return 0;
}